// typecheck.cpp
#include "typecheck.h"
#include <iostream>

// --- Public Methods ---

void TypeChecker::check(ASTNode* root) {
    if (!root) return;       //递归出口：NULL
    try {
        visit(root);
    } catch (const SemanticError& e) {
        std::cerr << "Semantic Error: " << e.what() << std::endl;        //输出错误日志
        exit(1);
    }
}

// --- Scope Management ---
//新增栈帧
void TypeChecker::enterScope() {
    scopes.emplace_back();
}

//弹出栈帧
void TypeChecker::exitScope() {
    if (!scopes.empty()) {
        scopes.pop_back();
    }
}

//声明变量
void TypeChecker::addSymbol(const SymbolInfo& symbol) {
    if (scopes.empty()) {
        throw SemanticError("Cannot add symbol, no scope active.");
    }
    auto& currentScope = scopes.back();
    if (currentScope.count(symbol.name)) {
        throw SemanticError("Redeclaration of symbol '" + symbol.name + "' in the same scope.");       //检测同一作用域内重定义
    }
    currentScope[symbol.name] = symbol;
}

//从符号表中查找
SymbolInfo* TypeChecker::lookupSymbol(const std::string& name) {
    for (auto it = scopes.rbegin(); it != scopes.rend(); ++it) {
        if (it->count(name)) {
            return &(*it)[name];
        }
    }
    return nullptr;
}


// --- AST Visitor Dispatcher ---

//自顶向下检查
void TypeChecker::visit(ASTNode* node) {
    if (node->type == "CompUnit") visitCompUnit(node);          //编译单元
    else if (node->type == "FuncDef") visitFuncDef(node);           //函数定义
    else {
        // Stmt 和 Expr 的分发由它们的父节点调用
        throw SemanticError("Unknown top-level node type: " + node->type);
    }
}


// --- Top-Level Visitor ---

//编译单元检查
void TypeChecker::visitCompUnit(ASTNode* node) {
    // 全局作用域
    enterScope();

    // 预扫描，将所有函数声明加入全局符号表，以支持先调用后定义（尽管我们的语言规范禁止）
    // 但更重要的是，这可以提前发现函数重定义
    for (const auto& child : node->children) {
        if (child->type == "FuncDef") {
            std::string returnType = child->children[0]->type;
            std::string funcName = child->children[1]->value;
            
            SymbolInfo funcSymbol;
            funcSymbol.name = funcName;
            funcSymbol.kind = SymbolKind::FUNCTION;
            funcSymbol.type = returnType;
            
            ASTNode* paramList = child->children[2].get();
            for (const auto& param : paramList->children) {
                // ToyC所有参数都是int
                funcSymbol.paramTypes.push_back("int");
            }
            addSymbol(funcSymbol);
        }
    }
    
    // 逐个访问函数定义
    for (const auto& child : node->children) {
        visit(child.get());
    }

    // 检查 main 函数
    SymbolInfo* mainFunc = lookupSymbol("main");
    if (!mainFunc) {
        throw SemanticError("Program entry point 'main' not found.");
    }
    if (mainFunc->kind != SymbolKind::FUNCTION) {
        throw SemanticError("'main' must be a function.");
    }
    if (mainFunc->type != "int") {
        throw SemanticError("'main' function must have 'int' return type.");
    }
    if (!mainFunc->paramTypes.empty()) {
        throw SemanticError("'main' function must not have parameters.");
    }
    
    exitScope();
}

//函数定义检查
void TypeChecker::visitFuncDef(ASTNode* node) {
    std::string returnType = node->children[0]->type;
    std::string funcName = node->children[1]->value;

    currentFunction = lookupSymbol(funcName);
    if (!currentFunction) {
         throw SemanticError("Internal error: Function '" + funcName + "' not found in pre-scan.");
    }
    
    if (lookupSymbol(funcName)->kind != SymbolKind::FUNCTION) {      //与非函数类型变量相撞
        throw SemanticError("'" + funcName + "' is already defined as a non-function symbol.");
    }

    enterScope(); // Function scope

    // 添加参数到函数作用域
    ASTNode* paramList = node->children[2].get();
    for (const auto& param : paramList->children) {
        std::string paramName = param->children[0]->value;
        addSymbol({paramName, SymbolKind::VARIABLE, "int"});
    }
    
    // 访问函数体
    ASTNode* body = node->children[3].get();
    bool allPathsReturn = visitBlock(body);        //检查是否函数每条路径都有返回值

    // 检查返回路径
    if (currentFunction->type == "int" && !allPathsReturn) {
        throw SemanticError("Not all control paths return a value in function '" + funcName + "'.");
    }

    exitScope();
    currentFunction = nullptr;
}


// --- Statement Visitors ---

//检查Statement
bool TypeChecker::visitStmt(ASTNode* node) {
    if (node->type == "Block") return visitBlock(node);
    if (node->type == "IfStmt" || node->type == "IfElseStmt") return visitIfStmt(node);
    if (node->type == "WhileStmt") return visitWhileStmt(node);
    if (node->type == "ReturnStmt") return visitReturnStmt(node);
    if (node->type == "BreakStmt") return visitBreakStmt(node);
    if (node->type == "ContinueStmt") return visitContinueStmt(node);
    if (node->type == "AssignStmt") return visitAssignStmt(node);
    if (node->type == "DeclAssignStmt") return visitDeclAssignStmt(node);
    if (node->type == "EmptyStmt") return visitEmptyStmt(node);
    
    // 表达式语句
    visitExpr(node);
    return false;
}

//检查函数体
bool TypeChecker::visitBlock(ASTNode* node) {
    enterScope();
    bool hasReturn = false;
    for (const auto& stmt : node->children) {
        // 如果已经有一条路径确定返回了，后续的代码就是不可达的，但我们仍然需要检查它们的语义
        if(visitStmt(stmt.get())) {
            hasReturn = true;
        }
    }
    exitScope();
    return hasReturn;
}

//if语句
bool TypeChecker::visitIfStmt(ASTNode* node) {
    // if (cond) then_stmt;
    // if (cond) then_stmt; else else_stmt;
    std::string condType = visitExpr(node->children[0].get());     //检查判断语句类型
    if (condType != "int") {
        throw SemanticError("If/While condition must be an integer (boolean) type, but got '" + condType + "'.");
    }
    
    //递归检查子语句
    bool thenReturns = visitStmt(node->children[1].get());
    if (node->type == "IfElseStmt") {
        bool elseReturns = visitStmt(node->children[2].get());
        return thenReturns && elseReturns;         //两条分支都有返回值才能保证各路径都有返回值
    }
    
    return false;   //对于不含else的if语句，如果条件为假则没有返回值，应返回false
}

//while语句
bool TypeChecker::visitWhileStmt(ASTNode* node) {
    std::string condType = visitExpr(node->children[0].get());      //检查判断条件类型
    if (condType != "int") {
        throw SemanticError("If/While condition must be an integer (boolean) type, but got '" + condType + "'.");
    }

    loopDepth++;
    visitStmt(node->children[1].get());  
    loopDepth--;

    return false;   //对于循环语句，如果一开始条件为假则没有返回值，应返回false
}

//return语句
bool TypeChecker::visitReturnStmt(ASTNode* node) {
    if (!currentFunction) {          //不在函数内但是却return了
        throw SemanticError("Return statement found outside of a function.");
    }
    
    if (node->children.empty()) { // return;(此时函数必须为void类型)
        if (currentFunction->type != "void") {
            throw SemanticError("Function '" + currentFunction->name + "' must return a value of type 'int'.");
        }
    } else { // return expr; 
        /*
        if (currentFunction->type == "void") {
            throw SemanticError("Void function '" + currentFunction->name + "' cannot return a value.");
        }
        */
        std::string returnExprType = visitExpr(node->children[0].get());         //检查return表达式的类型
        if (returnExprType != currentFunction->type) {          //表达式类型与函数返回类型不匹配，应报错
            throw SemanticError("Type mismatch in return statement. Expected '" + currentFunction->type + "' but got '" + returnExprType + "'.");
        }
    }
    return true;  //递归出口：表示有return path
}

//break语句
bool TypeChecker::visitBreakStmt(ASTNode* node) {
    if (loopDepth == 0) {        //检查循环深度，必须在循环中才能break/continue
        throw SemanticError("'break' statement not in loop.");
    }
    return false;
}

//continue语句
bool TypeChecker::visitContinueStmt(ASTNode* node) {
    if (loopDepth == 0) {     //检查循环深度，必须在循环中才能break/continue
        throw SemanticError("'continue' statement not in loop.");
    }
    return false;
}

//赋值语句
bool TypeChecker::visitAssignStmt(ASTNode* node) {
    std::string varName = node->children[0]->value;
    SymbolInfo* symbol = lookupSymbol(varName);
    if (!symbol) {        //未声明的变量
        throw SemanticError("Undeclared variable '" + varName + "'.");
    } 
    if (symbol->kind == SymbolKind::FUNCTION) {        //检查到左值是函数
        throw SemanticError("Cannot assign to a function '" + varName + "'.");
    }
    
    std::string rhsType = visitExpr(node->children[1].get());
    if (symbol->type != rhsType) {         //类型不匹配
        throw SemanticError("Type mismatch in assignment to '" + varName + "'. Expected '" + symbol->type + "' but got '" + rhsType + "'.");
    }
    return false;
}

//声明+赋值语句
bool TypeChecker::visitDeclAssignStmt(ASTNode* node) {
    // std::string typeName = node->children[0]->type; // Always "int" in ToyC
    std::string varName = node->children[1]->value;
    
    addSymbol({varName, SymbolKind::VARIABLE, "int"});
    
    std::string rhsType = visitExpr(node->children[2].get());
    if (rhsType != "int") {         //void类型不能同时声明+赋值
        throw SemanticError("Type mismatch in declaration of '" + varName + "'. Expected 'int' but got '" + rhsType + "'.");
    }
    return false;
}

//空语句
bool TypeChecker::visitEmptyStmt(ASTNode* node) {
    return false;
}


// --- Expression Visitors ---

//检查表达式，返回表达式的类型
std::string TypeChecker::visitExpr(ASTNode* node) {
    if (node->type == "LOrExpr" || node->type == "LAndExpr" ||
        node->type == "RelExpr" || node->type == "AddExpr" ||
        node->type == "MulExpr") {
        return visitBinaryExpr(node);
    }
    if (node->type == "UnaryExpr") return visitUnaryExpr(node);
    if (node->type == "VarRef") return visitVarRef(node);
    if (node->type == "NumLiteral") return visitNumLiteral(node);
    if (node->type == "FuncCall") return visitFuncCall(node);
    
    // For cases like (Expr)
    if (!node->children.empty()) {
        return visitExpr(node->children[0].get());
    }

    throw SemanticError("Unsupported expression type: " + node->type);
}

//二元表达式
std::string TypeChecker::visitBinaryExpr(ASTNode* node) {
    std::string leftType = visitExpr(node->children[0].get());
    std::string rightType = visitExpr(node->children[1].get());
    
    //两边均需要为int类型
    if (leftType != "int" || rightType != "int") {
        throw SemanticError("Operands of binary operator '" + node->value + "' must be of type 'int'.");
    }
    
    // All binary expressions in ToyC return int (or boolean, which is int)
    return "int";
}

//一元表达式
std::string TypeChecker::visitUnaryExpr(ASTNode* node) {
    std::string operandType = visitExpr(node->children[0].get());
    if (operandType != "int") {
        throw SemanticError("Operand of unary operator '" + node->value + "' must be of type 'int'.");
    }
    return "int";
}

//值（Var）
std::string TypeChecker::visitVarRef(ASTNode* node) {
    SymbolInfo* symbol = lookupSymbol(node->value);
    if (!symbol) {       //未声明的变量
        throw SemanticError("Undeclared variable '" + node->value + "'.");
    }           
    if (symbol->kind == SymbolKind::FUNCTION) {         //函数名不能出现在表达式中
        throw SemanticError("Cannot use function '" + node->value + "' as a value.");
    }
    return symbol->type;
}

//数字
std::string TypeChecker::visitNumLiteral(ASTNode* node) {
    return "int";
}

//函数调用
std::string TypeChecker::visitFuncCall(ASTNode* node) {
    std::string funcName = node->value;
    SymbolInfo* symbol = lookupSymbol(funcName);
    
    if (!symbol) {
        // 由于我们的语言规范要求先定义后调用，所以这里找不到就是错误
        throw SemanticError("Call to undeclared function '" + funcName + "'.");
    } 
    if (symbol->kind != SymbolKind::FUNCTION) {      //并未函数类型
        throw SemanticError("'" + funcName + "' is not a function and cannot be called.");
    }

    // 检查参数数量
    size_t expected_args = symbol->paramTypes.size();
    size_t actual_args = node->children.size();
    if (expected_args != actual_args) {
        throw SemanticError("Function '" + funcName + "' expects " + std::to_string(expected_args) + 
                            " arguments, but " + std::to_string(actual_args) + " were provided.");
    }

    // 检查参数类型
    for (size_t i = 0; i < actual_args; ++i) {
        std::string argType = visitExpr(node->children[i].get());
        if (argType != symbol->paramTypes[i]) {
            throw SemanticError("Type mismatch for argument " + std::to_string(i+1) + " in call to '" + 
                                funcName + "'. Expected '" + symbol->paramTypes[i] + "' but got '" + argType + "'.");
        }
        // 检查void返回的函数调用是否被用作值
        if (argType == "void") {
             throw SemanticError("A function with 'void' return type cannot be used as an argument.");
        }
    }
    
    return symbol->type;
}
